#!/usr/bin/env python

import xlh
import sys, getopt, os
from random import Random

# Define some general functions


def exitOnError(errmsg, exit):
	print 'The following error(s) occurred:\n'
	print errmsg
	sys.exit(exit)
def usage():
	print """Usage:

""",sys.argv[0],""" [-h] -d DATABASE [-f FILE | -a ACCESSIONS] -m MISSED [-e ENZYME] [-o OUTPUT] [-l LENGTH] [-x]

 -h:	Print this help and exit
 -d:	Database used for previous MS search
 -f:	File containing accession numbers
 -a:	Comma separated accession numbers
 -m:	Number of missed cleavages used in previous MS search [Default: 2]
 -e:	Enzyme [Default: Trypsin]
 -o:	Name of output database [Optional]
 -l:	Minumum peptide length [Default: 3]
 -L:	Maximum peptide length [Default: 15]
 -A:	Use all peptides (not only xlinking ones) [Default: False]
 -x:	Avoid intrachain xlinking peptides (whenever possibile) [Default: False]
	"""
	sys.exit(1)

def readAccessionFile(fileName):
	accs = []
	fh = open(fileName, 'r')
	while 1:
		l = fh.readline()
		if not l:
			break
		accs.append(l.strip())
	fh.close()
	return accs
	

MinLen = 3
MaxLen = 15
DbIn = ''
Accs = []
Missed = 2
DbOut = ''
Function = 0
Enzyme = 'Trypsin'
UseAll = False
AvoidIntra = False

try:
	optlist, args = getopt.getopt(sys.argv[1:], 'd:f:a:m:o:e:l:L:Ax')
except getopt.GetoptError:
	usage()

if not optlist:
	usage()

for o, a in optlist:
	if o == '-h':
		usage()
	if o == '-d':
		DbIn = a
	if o == '-o':
		DbOut = a
	if o == '-m':
		Missed = int(a)
	if o == '-l':
		MinLen = int(a)
	if o == '-L':
		MaxLen = int(a)
	if o == '-f':
		Accs = readAccessionFile(a)
	if o == '-a':
		Accs = a.split(',')
	if o == '-e':
		Enzyme = a.strip()
	if o == '-A':
		UseAll = True
	if o == '-x':
		AvoidIntra = True
		
# do some checks
if not DbIn:
	exitOnError('Missing database name', 2)
if not Accs:
	exitOnError('No accession numbers defined', 2)
if not DbOut:
	DbOut = 'db' + str(Random().getrandbits(64)) 

FastaOut = DbOut + '.fasta'
DBtitle = 'Used ' + Enzyme + ' with ' + str(Missed) + ' miscleavages'

# read proteins from database

SingleProteins = []
for acc in Accs:
	SingleProteins.append(xlh.Protein(xlh.readEntryFromNCBI(DbIn, acc)))
	SingleProteins[-1].cleave(Enzyme, Missed)

	
# in this case we should create a pool of peptides
# and concatenate them ABA, BCB, CDC, DAD etc.
# each concatenation is a single fasta entry
t = {}
entries = []
for pr in SingleProteins:
	if UseAll:
		peptides = pr.peptideSequences()
	else:
		peptides = filter(lambda x: xlh.Peptide(x).isNTerminal(pr), pr.peptideSequences())
		peptides += filter(lambda x: 'K' in x[:-1], pr.peptideSequences())
	for pt in peptides:
		try:
			t[pt].append(pr.accession())
		except KeyError:
			t[pt] = [pr.accession()]
peps = filter(lambda x: len(x) >= MinLen and len(x) <= MaxLen, t.keys())
rough = len(peps) * len(peps) / 2
dec = rough / 20

sys.stdout.write("Got " + str(len(peps)) + " peptides\n")

sys.stdout.write("Writing " + FastaOut)
sys.stdout.flush()

linewrap = 80
thisDir = os.getcwd()
if '/' not in FastaOut:
	# if no directory is given, assume that is the
	# one specified in external.ini
	os.chdir(xlh.External.get("blastdb", "path"))

fh = open(FastaOut, 'w')

n = 0
for x in range(len(peps) -1):
	for y in range(x, len(peps)):
		if AvoidIntra and len(filter(lambda pepCheck: pepCheck in t[peps[y]], t[peps[x]])):
			continue
		if (n) % dec == 0:
			sys.stderr.write('.')
			sys.stderr.flush()
		acc = 'xpep' + str(x) + 'x' + str(y)
		desc = 'xlh|' + acc + '| ' + peps[x] + ':' + ','.join(t[peps[x]]) + ' ' + peps[y] + ':'  + ','.join(t[peps[y]])
		seq = peps[x] + peps[y] + peps[x]
		fh.write('>' + desc + '\n')
		for i in range(len(seq) / linewrap):
			fh.write(seq[i * linewrap:i * linewrap + linewrap] + '\n')
		r = len(seq) % linewrap
		fh.write(seq[-r:] + '\n')
		n += 1
sys.stdout.flush()
sys.stdout.write(" OK\n")

fh.close()
os.chdir(thisDir)

sys.stdout.write("Formatting for omssa search...")
sys.stdout.flush()
xlh.writeNCBI(FastaOut, DbOut, DBtitle)
sys.stdout.write(" OK\n")
